import gc
import io
import re
from typing import Dict, Any, List
from urllib.parse import urlparse  # 用于解析URL获取路径和文件名
import requests  # 用于网络请求获取在线文档
import logging  # 用于日志记录
from pathlib import Path  # 用于处理文件路径

import fitz  # PyMuPDF库，用于解析PDF文件
from bs4 import BeautifulSoup  # 用于解析HTML内容
from langchain_core.tools import Tool  # 用于创建LangChain工具
from tqdm import tqdm  # 用于显示进度条

from docx import Document  # 用于解析DOCX文件
from docx.opc.exceptions import PackageNotFoundError  # DOCX文件损坏时的异常

from core.exception import ApiException  # 自定义API异常类

# 初始化日志记录器
logger = logging.getLogger(__name__)


class DocumentTools:
    """文档处理工具类

    提供在线文档的获取、解析和内容提取功能，支持多种格式（HTML/PDF/TXT/MD/DOCX），
    并能将解析结果封装为统一格式。同时提供工具注册方法，用于LangChain集成。
    """

    def __init__(self):
        """初始化文档工具类

        创建全局requests会话，配置连接池以提高多个请求的效率和性能。
        连接池设置：10个连接，最大10个并发，不阻塞等待。
        """
        # 创建全局Session复用连接，减少TCP握手开销
        self.session = requests.Session()
        # 配置HTTP连接池，优化网络请求性能
        adapter = requests.adapters.HTTPAdapter(
            pool_connections=10,  # 连接池中的连接数
            pool_maxsize=10,      # 连接池的最大大小
            pool_block=False      # 当连接池满时不阻塞，直接创建新连接
        )
        # 为HTTP和HTTPS协议挂载适配器
        self.session.mount('http://', adapter)
        self.session.mount('https://', adapter)

    @staticmethod
    def fetch_document(url: str) -> Dict[str, Any]:
        """获取并解析在线文档

        根据文档的Content-Type和URL后缀判断文档类型，调用对应解析方法，
        提取文档标题和内容，并返回统一格式的结果。

        Args:
            url: 在线文档的URL地址

        Returns:
            包含文档信息的字典，结构为：
            {
                "title": 文档标题,
                "content": 提取的文档内容,
                "type": 文档类型(pdf/html/txt/md/docx/text)
            }

        Raises:
            ApiException: 当文档获取或解析失败时抛出
        """
        try:
            # 发送GET请求获取文档，stream=True表示流式下载，适合大文件
            response = DocumentTools().session.get(url, stream=True, timeout=30)
            # 检查请求是否成功（状态码200-299），否则抛出异常
            response.raise_for_status()

            # 获取响应的Content-Type，用于判断文档类型
            content_type = response.headers.get('content-type', "").lower()
            # 从URL路径中提取文件后缀，辅助判断文档类型
            url_suffix = Path(urlparse(url).path).suffix.lower()

            # 处理DOCX格式：根据MIME类型或文件后缀判断
            if ("application/vnd.openxmlformats-officedocument.wordprocessingml.document" in content_type) or (
                    url_suffix == ".docx"):
                logger.info(f"解析DOCX文档，URL: {url}")
                return DocumentTools._parse_docx(response)

            # 处理Markdown格式：根据MIME类型或文件后缀判断
            elif ("text/markdown" in content_type) or (url_suffix in [".md", ".markdown"]):
                logger.info(f"解析Markdown文档，URL: {url}")
                return DocumentTools._parse_md(response)

            # 处理TXT格式：根据MIME类型或文件后缀判断
            elif ("text/plain" in content_type) or (url_suffix == ".txt"):
                logger.info(f"解析TXT文档，URL: {url}")
                return DocumentTools._parse_txt(response, url)

            # 处理HTML格式：根据MIME类型判断
            elif "text/html" in content_type:
                logger.info(f"解析HTML文档，URL: {url}")
                return DocumentTools._parse_html(response)

            # 处理PDF格式：根据MIME类型判断
            elif "application/pdf" in content_type:
                logger.info(f"解析PDF文档，URL: {url}")
                return DocumentTools._parse_pdf(response)

            # 未知格式按纯文本处理
            else:
                logger.warning(f"未知类型，按纯文本处理，URL: {url}")
                # 使用URL路径中的文件名作为标题，若无则使用默认值
                title = Path(urlparse(url).path).name or "未知文件名"
                return {
                    "title": title,
                    "content": DocumentTools._clean_text(response.text),
                    "type": "text"
                }

        except Exception as e:
            # 记录错误日志并抛出自定义异常
            logger.error(f"文档处理失败，URL: {url}，错误: {str(e)}", exc_info=True)
            raise ApiException(msg=f"文档处理失败：{str(e)}") from e

    # 以下为各类文档的解析方法（私有静态方法）
    @staticmethod
    def _parse_html(response: requests.Response) -> Dict[str, Any]:
        """解析HTML文档

        使用BeautifulSoup提取HTML中的标题和文本内容，标题优先从<title>标签获取，
        内容提取所有文本并进行清洗。

        Args:
            response: 包含HTML内容的响应对象

        Returns:
            包含HTML文档信息的字典
        """
        # 使用lxml解析器解析HTML内容
        soup = BeautifulSoup(response.text, "lxml")
        # 提取标题：优先使用<title>标签内容，若无则使用默认值
        title = (soup.title.string.strip() if (soup.title and soup.title.string) else "未命名文档")
        # 提取所有文本内容，使用换行符分隔，去除首尾空白
        content = soup.get_text(separator="\n", strip=True)
        # 清洗文本内容
        content = DocumentTools._clean_text(content)
        return {"title": title, "content": content, "type": "html"}

    @staticmethod
    def _parse_pdf(response: requests.Response) -> Dict[str, Any]:
        """使用PyMuPDF解析PDF文档

        流式下载PDF文件到内存缓冲区，使用PyMuPDF读取并提取每一页的文本内容，
        标题优先从PDF元数据获取，若无则使用URL中的文件名。

        Args:
            response: 包含PDF内容的响应对象

        Returns:
            包含PDF文档信息的字典

        Raises:
            ApiException: 当PDF解析失败时抛出
        """
        # 创建内存缓冲区存储PDF内容
        buffer = io.BytesIO()
        try:
            # 获取文件总大小，用于进度条显示
            total_size = int(response.headers.get("content-length", 0))
            # 显示下载进度条
            with tqdm(total=total_size, unit="B", unit_scale=True, desc="下载PDF", mininterval=0.5) as pbar:
                # 分块读取响应内容并写入缓冲区
                for chunk in response.iter_content(chunk_size=32768):
                    if chunk:  # 过滤空块
                        buffer.write(chunk)
                        pbar.update(len(chunk))  # 更新进度条

            # 将缓冲区指针移到开头，准备读取
            buffer.seek(0)
            # 使用PyMuPDF打开内存中的PDF文件
            with fitz.open(stream=buffer, filetype="pdf") as doc:
                # 提取标题：优先使用PDF元数据，若无则使用URL中的文件名
                title = doc.metadata.get("title", "").strip() or Path(urlparse(response.url).path).name or "未命名PDF"
                content = []
                # 显示解析进度条
                with tqdm(total=len(doc), desc="解析PDF", mininterval=0.5) as pbar:
                    # 逐页提取文本
                    for page in doc:
                        content.append(page.get_text().strip())
                        pbar.update(1)  # 更新进度条
                # 将所有页的文本合并
                content = "\n".join(content)
        # 清洗文本内容
            content = DocumentTools._clean_text(content)
            return {"title": title, "content": content, "type": "pdf"}

        except Exception as e:
            # 捕获解析异常并抛出自定义异常
            raise ApiException(msg=f"PDF解析失败：{str(e)}") from e
        finally:
            # 确保缓冲区关闭，释放资源
            buffer.close()

    @staticmethod
    def _parse_txt(response: requests.Response, url: str) -> Dict[str, Any]:
        """解析TXT文档

        流式读取TXT文件内容，逐行处理并合并，标题使用URL中的文件名。

        Args:
            response: 包含TXT内容的响应对象
            url: 文档URL，用于提取标题

        Returns:
            包含TXT文档信息的字典

        Raises:
            ApiException: 当TXT解析失败时抛出
        """
        try:
            # 使用URL中的文件名作为标题
            title = Path(urlparse(url).path).name or "未命名TXT文件"
            content = []
            # 流式读取文本行（自动解码为Unicode）
            for line in response.iter_lines(decode_unicode=True):
                if line:  # 过滤空行
                    content.append(line.strip())
            # 合并所有行文本
            content = "\n".join(content)
            # 清洗文本内容
            content = DocumentTools._clean_text(content)
            return {"title": title, "content": content, "type": "txt"}
        except Exception as e:
            # 捕获解析异常并抛出自定义异常
            raise ApiException(msg=f"TXT解析失败：{str(e)}") from e

    @staticmethod
    def _parse_md(response: requests.Response) -> Dict[str, Any]:
        """解析Markdown文档

        提取Markdown中的一级标题（# 标题）作为文档标题，若无则使用URL中的文件名，
        内容保留原始Markdown格式并进行清洗。

        Args:
            response: 包含Markdown内容的响应对象

        Returns:
            包含Markdown文档信息的字典

        Raises:
            ApiException: 当Markdown解析失败时抛出
        """
        try:
            # 读取Markdown文本内容
            md_text = response.text

            # 正则匹配一级标题（# 标题）
            title_match = re.search(r'^# (.+)$', md_text, re.MULTILINE)
            # 提取标题：优先使用一级标题，若无则使用URL中的文件名
            title = title_match.group(1).strip() if title_match else Path(urlparse(response.url).path).name or "未命名MD文件"

            # 清洗文本内容
            content = DocumentTools._clean_text(md_text)
            return {"title": title, "content": content, "type": "md"}
        except Exception as e:
            # 捕获解析异常并抛出自定义异常
            raise ApiException(msg=f"MD解析失败：{str(e)}") from e

    @staticmethod
    def _parse_docx(response: requests.Response) -> Dict[str, Any]:
        """解析DOCX文档

        流式下载DOCX文件到内存缓冲区，使用python-docx库读取内容，
        标题优先从文档属性获取，若无则使用第一个一级标题，否则使用URL中的文件名。

        Args:
            response: 包含DOCX内容的响应对象

        Returns:
            包含DOCX文档信息的字典

        Raises:
            ApiException: 当DOCX解析失败或文件损坏时抛出
        """
        # 创建内存缓冲区存储DOCX内容
        buffer = io.BytesIO()
        try:
            # 获取文件总大小，用于进度条显示
            total_size = int(response.headers.get("content-length", 0))
            # 显示下载进度条
            with tqdm(total=total_size, unit="B", unit_scale=True, desc="下载DOCX", mininterval=0.5) as pbar:
                # 分块读取响应内容并写入缓冲区
                for chunk in response.iter_content(chunk_size=32768):
                    if chunk:  # 过滤空块
                        buffer.write(chunk)
                        pbar.update(len(chunk))  # 更新进度条

            # 将缓冲区指针移到开头，准备读取
            buffer.seek(0)
            # 使用python-docx打开内存中的DOCX文件
            doc = Document(buffer)

            # 提取标题：优先使用文档属性中的标题
            title = "未命名DOCX文件"
            if doc.core_properties.title:
                title = doc.core_properties.title.strip()
            else:
                # 若无文档属性标题，则查找第一个一级标题
                for paragraph in doc.paragraphs:
                    if paragraph.style.name == "Heading 1" and paragraph.text.strip():
                        title = paragraph.text.strip()
                        break
                # 若仍未找到标题，则使用URL中的文件名
                if title == "未命名DOCX文件":
                    title = Path(urlparse(response.url).path).name or title

            # 提取内容
            content = []
            total_paragraphs = len(doc.paragraphs)
            # 显示解析进度条
            with tqdm(total=total_paragraphs, desc="解析DOCX", mininterval=0.5) as pbar:
                for idx, paragraph in enumerate(doc.paragraphs):
                    para_text = paragraph.text.strip()
                    if para_text:  # 过滤空段落
                        content.append(para_text)
                    # 每处理50个段落进行一次垃圾回收，优化内存使用
                    if (idx + 1) % 50 == 0:
                        gc.collect()
                    pbar.update(1)  # 更新进度条
            # 合并所有段落文本
            content = "\n".join(content)
            # 清洗文本内容
            content = DocumentTools._clean_text(content)
            return {"title": title, "content": content, "type": "docx"}

        except PackageNotFoundError:
            # 处理DOCX文件损坏的情况
            raise ApiException(msg="DOCX文件损坏或格式不合法") from None
        except Exception as e:
            # 捕获其他解析异常并抛出自定义异常
            raise ApiException(msg=f"DOCX解析失败：{str(e)}") from e
        finally:
            # 确保缓冲区关闭，释放资源
            buffer.close()

    @staticmethod
    def _clean_text(text: str) -> str:
        """清理文本内容

        标准化文本格式：将连续3个及以上换行符替换为2个，
        将连续空格或制表符替换为单个空格，并去除首尾空白。

        Args:
            text: 需要清理的原始文本

        Returns:
            清理后的文本
        """
        # 替换连续3个及以上换行符为2个换行符
        text = re.sub(r'\n{3,}', '\n\n', text)
        # 替换连续空格或制表符为单个空格
        text = re.sub(r'[ \t]+', ' ', text)
        # 去除首尾空白
        return text.strip()

    @staticmethod
    def create_tools() -> List[Tool]:
        """创建LangChain工具列表

        将文档获取解析功能封装为LangChain的Tool对象，用于在链中调用。

        Returns:
            包含文档处理工具的列表
        """
        fetch_doc_tool = Tool(
            name="fetch_document",  # 工具名称
            func=DocumentTools.fetch_document,  # 工具执行函数
            description="""
            用于获取并解析在线文档（支持HTML/PDF/TXT/MD/DOCX），提取标题和正文。
            参数：有效的URL字符串（如https://example.com/test.pdf）。
            注意：文档需可公开访问，不支持加密文档。
            """  # 工具描述，用于提示词生成
        )
        return [fetch_doc_tool]